#include <core/main.h>

void NeighborhoodFeatureSet::buildFromReplays(const ReplayCollection &replays, int num, int size, int gridScale)
{
    features.clear();
    _size = size;
    NeighborhoodCollection neighborhoods;
    for (auto it = replays.Replays().begin(); it != replays.Replays().end(); ++it)
    {
        Replay *replay = *it;
        cout << "Looking at a replay from game " << replay->SourceGame()->Name() << endl;
        int frameCount = 0;
        for (auto fit = replay->Frames().begin(); fit != replay->Frames().end(); ++fit)
        {
            if (frameCount % 1000 == 0) cout << "Looking at frame " << frameCount << endl;
            frameCount++;
            const Frame &frame = *fit;
            FrameGrid frameGrid(frame, gridScale);
            neighborhoods.addAll(frameGrid, size);
        }
    }
    neighborhoods.getMostFrequent(num, features);
}

void NeighborhoodFeatureSet::saveToFile(const string &filename) const
{
    OutputDataStream out;
    out << *this;
    out.SaveToFile(filename);
}

void NeighborhoodFeatureSet::loadFromFile(const string &filename)
{
    InputDataStream is;
    is.LoadFromFile(filename);
    is >> *this;
}

void NeighborhoodFeatureSet::exportPNGs(const SpriteCollection &spriteCollection, const string &dirname) const
{
    int figNum = 1;
    for (auto it = features.begin(); it != features.end(); ++it) {
        stringstream figName;
        figName << dirname << figNum << ".png";
        Bitmap bmp = it->getBitmap(spriteCollection);
        bmp.SavePNG(figName.str());
        figNum++;
    }
}

void NeighborhoodFeatureSet::encodeFrame(const FrameGrid &frame, vector<double> &out) const
{
    map<Neighborhood, int> countsInFrame;
    for (int y = 0; y + _size <= frame.height(); y++) {
        for (int x = 0; x + _size <= frame.width(); x++) {
            Neighborhood neigh = frame.getNeighborhood(x, y, _size);
            if (features.count(neigh)) {
                countsInFrame[neigh]++;
            }
        }
    }

    out.clear();
    for (auto it = features.begin(); it != features.end(); ++it) {
        if (countsInFrame.count(*it)) {
            out.push_back(countsInFrame[*it]);
        } else {
            out.push_back(0);
        }
    }
}

OutputDataStream & operator<< (OutputDataStream &os, const NeighborhoodFeatureSet &nfs)
{
    os << nfs._size;
    os << nfs.features.size();
    for (auto it = nfs.features.begin(); it != nfs.features.end(); ++it) {
        os << *it;
    }
    return os;
}

InputDataStream & operator>> (InputDataStream &is, NeighborhoodFeatureSet &nfs)
{
    is >> nfs._size;
    int numFeatures;
    is >> numFeatures;
    for (int i = 0; i < numFeatures; i++) {
        Neighborhood neigh;
        is >> neigh;
        nfs.features.insert(neigh);
    }
    return is;
}